--[[
令牌桶消费脚本 - 使用JSON格式存储
KEYS[1]: bucket key
ARGV[1]: capacity
ARGV[2]: refillRate  
ARGV[3]: requestedTokens
ARGV[4]: currentTime
ARGV[5]: expireSeconds

返回: {allowed(0|1), remainingTokens, waitTime}
--]]

local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local refillRate = tonumber(ARGV[2])
local requestedTokens = tonumber(ARGV[3])
local now = tonumber(ARGV[4])
local expireSeconds = tonumber(ARGV[5])

-- 获取当前桶状态 (JSON格式)
local bucketData = redis.call('GET', key)
local tokens, lastRefillTime

if bucketData then
    -- 解析JSON数据
    local bucket = cjson.decode(bucketData)
    tokens = bucket.tokens or capacity
    lastRefillTime = bucket.lastRefillTime or now
else
    -- 初始化新桶
    tokens = capacity
    lastRefillTime = now
end

-- 补充令牌
if now > lastRefillTime then
    local timeDiff = now - lastRefillTime
    local tokensToAdd = math.floor(timeDiff * refillRate)
    tokens = math.min(capacity, tokens + tokensToAdd)
    lastRefillTime = now
end

-- 尝试消费令牌
local allowed = 0
local waitTime = 0

if tokens >= requestedTokens then
    tokens = tokens - requestedTokens
    allowed = 1
else
    -- 计算等待时间
    local deficit = requestedTokens - tokens
    waitTime = math.ceil(deficit / refillRate)
end

-- 保存桶状态 (JSON格式)
local bucketState = {
    capacity = capacity,
    tokens = tokens,
    refillRate = refillRate,
    lastRefillTime = lastRefillTime,
    updatedAt = now
}

redis.call('SET', key, cjson.encode(bucketState))
redis.call('EXPIRE', key, expireSeconds)

return {allowed, tokens, waitTime} 